#include "net/arp.h"
#include "drivers/pcap_adapter.h"
#include <string.h>
#include <stdlib.h>
#include <stdio.h>

#ifdef _WIN32
/* Winsock (provide htons/ntohs/htonl/ntohl) */
#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#include <winsock2.h>
#include <Ws2tcpip.h>
#include <windows.h>
#else
#include <arpa/inet.h>
#include <sys/time.h>
#endif

#define ARP_CACHE_SIZE 32
#define ARP_QUEUE_MAX 32

#pragma pack(push,1)
typedef struct {
    uint16_t hw_type;
    uint16_t proto_type;
    uint8_t hw_len;
    uint8_t proto_len;
    uint16_t opcode;
    uint8_t sha[6];
    uint32_t spa;
    uint8_t tha[6];
    uint32_t tpa;
} arp_pkt_t;
#pragma pack(pop)

typedef struct {
    uint32_t ip;
    uint8_t mac[6];
    uint32_t ts; // timestamp
    int valid;
} arp_entry_t;

typedef struct {
    uint8_t* frame;
    size_t len;
    uint32_t dst_ip;
} arp_queued_frame_t;

static void* g_netif = NULL;
static uint8_t g_src_mac[6];
static uint32_t g_src_ip = 0;
static arp_entry_t g_cache[ARP_CACHE_SIZE];
static arp_queued_frame_t g_queue[ARP_QUEUE_MAX];

static uint32_t get_time_sec() {
#ifdef _WIN32
    return (uint32_t)(GetTickCount() / 1000);
#else
    struct timeval tv; gettimeofday(&tv, NULL); return (uint32_t)tv.tv_sec;
#endif
}

int arp_init(void* netif_handle, const uint8_t src_mac[6], uint32_t src_ip) {
    g_netif = netif_handle;
    memcpy(g_src_mac, src_mac, 6);
    g_src_ip = src_ip;
    memset(g_cache,0,sizeof(g_cache));
    memset(g_queue,0,sizeof(g_queue));
    return 0;
}

static void send_arp_request(uint32_t target_ip) {
    // build ethernet + arp request
    size_t eth = 14;
    size_t arp_sz = sizeof(arp_pkt_t);
    size_t tot = eth + arp_sz;
    uint8_t* buf = (uint8_t*)malloc(tot);
    if (!buf) return;
    // eth dst = ff:ff:ff:ff:ff:ff
    memset(buf,0xff,6);
    memcpy(buf+6,g_src_mac,6);
    *(uint16_t*)(buf+12) = htons(0x0806);
    arp_pkt_t* a = (arp_pkt_t*)(buf+eth);
    a->hw_type = htons(1);
    a->proto_type = htons(0x0800);
    a->hw_len = 6;
    a->proto_len = 4;
    a->opcode = htons(1); // request
    memcpy(a->sha, g_src_mac, 6);
    a->spa = htonl(g_src_ip);
    memset(a->tha,0,6);
    a->tpa = htonl(target_ip);
    pcap_adapter_send(g_netif, buf, tot);
    free(buf);
}

int arp_resolve(uint32_t ip, uint8_t out_mac[6]) {
    for (int i=0;i<ARP_CACHE_SIZE;i++) {
        if (g_cache[i].valid && g_cache[i].ip == ip) {
            memcpy(out_mac, g_cache[i].mac, 6);
            return 0;
        }
    }
    return -1;
}

int arp_queue_frame(uint32_t dst_ip, const uint8_t* frame, size_t len) {
    for (int i=0;i<ARP_QUEUE_MAX;i++) {
        if (g_queue[i].frame == NULL) {
            g_queue[i].frame = (uint8_t*)malloc(len);
            if (!g_queue[i].frame) return -1;
            memcpy(g_queue[i].frame, frame, len);
            g_queue[i].len = len;
            g_queue[i].dst_ip = dst_ip;
            // send ARP request immediately
            send_arp_request(dst_ip);
            return 0;
        }
    }
    return -1;
}

int arp_input(const uint8_t* pkt, size_t len) {
    if (len < 14 + sizeof(arp_pkt_t)) return -1;
    uint16_t ethertype = ntohs(*(uint16_t*)(pkt+12));
    if (ethertype != 0x0806) return 0;
    const arp_pkt_t* a = (const arp_pkt_t*)(pkt+14);
    uint16_t opcode = ntohs(a->opcode);
    uint32_t spa = ntohl(a->spa);
    uint8_t* sha = (uint8_t*)a->sha;
    // update cache
    int found = -1;
    for (int i=0;i<ARP_CACHE_SIZE;i++) {
        if (g_cache[i].valid && g_cache[i].ip == spa) { found=i; break; }
    }
    if (found == -1) {
        for (int i=0;i<ARP_CACHE_SIZE;i++) if (!g_cache[i].valid) { found=i; break; }
        if (found == -1) { found = rand() % ARP_CACHE_SIZE; }
    }
    g_cache[found].ip = spa;
    memcpy(g_cache[found].mac, sha, 6);
    g_cache[found].ts = get_time_sec();
    g_cache[found].valid = 1;
    // if request is reply for queued frames, send them
    for (int i=0;i<ARP_QUEUE_MAX;i++) {
        if (g_queue[i].frame && g_queue[i].dst_ip == spa) {
            // patch dst MAC in queued frame (assume eth dst at offset 0)
            memcpy(g_queue[i].frame, sha, 6);
            pcap_adapter_send(g_netif, g_queue[i].frame, g_queue[i].len);
            free(g_queue[i].frame);
            g_queue[i].frame = NULL;
            g_queue[i].len = 0;
        }
    }
    return 0;
}

void arp_tick(void) {
    uint32_t now = get_time_sec();
    for (int i=0;i<ARP_CACHE_SIZE;i++) {
        if (g_cache[i].valid && (now - g_cache[i].ts) > 120) {
            g_cache[i].valid = 0;
        }
    }
    // Optionally re-send ARP requests for queued frames (not implemented full)
}
